# -*- coding: utf8 -*-

import sys
from collections import defaultdict, namedtuple

import config
# noinspection PyUnresolvedReferences
from config import AlignDatabaseLine  # must be for unpickle object from dump
from utils import profile, load_object
from func_util import head, last


ExonInfo = namedtuple('ExonInfo', ['start', 'minus_end', 'strand', 'name'])

EXON_DISTANCE_INCREASE_MULTIPLIER = 1.5
MIN_EXON_LENGTH = 50
MIN_GENE_EXONS_COUNT = 3


def main():
    gene_struct = load_genes_struct()
    align_db = load_object(config.EXON_ALIGN_DUMP_FILENAME)
    segment_length = load_object(config.SEGMENT_LENGTH_DUMP_FILENAME)

    aligned = not_aligned = 0
    for number, (gene_name, exons) in enumerate(gene_struct.items()):
        # filter all genes with too few exons
        if len(exons) < MIN_GENE_EXONS_COUNT:
            continue
        gene_aligned = process_gene(align_db, segment_length, sorted(exons), gene_name)
        if gene_aligned:
            aligned += gene_aligned
        else:
            not_aligned += 1

    print('stat: aligned {}, not aligned {}'.format(aligned, not_aligned))


@profile
def load_genes_struct():
    """ gene -> [exon :: (start, end, strand, name)] """
    sys.stderr.write('Load struct...\n')
    gene_struct = defaultdict(set)
    with open(config.STRUCTURE_FILENAME) as handle:
        handle.readline()
        for line in handle.readlines():
            gene_name, exon_name, *numbers = line.split()
            start, end, strand = map(int, numbers)
            if end - start > MIN_EXON_LENGTH:
                gene_struct[gene_name].add(ExonInfo(start, -end, strand, exon_name))

    filtered_gene_struct = {}
    for gene_name, exons in gene_struct.items():
        exons = sorted(exons)
        filtered = [head(exons)]
        for i in range(1, len(exons)):
            if calc_exon_distance(last(filtered), exons[i]) >= 0:
                filtered.append(exons[i])
        if len(filtered) < 2:
            continue
        filtered_gene_struct[gene_name] = filtered
    return filtered_gene_struct


def process_gene(align_db, segment_length_table, sorted_gene_exons, gene_name):
    start_exon = head(sorted_gene_exons)
    start_exon_align_lines = sorted(align_db[start_exon.name])
    aligned_gene = 0
    for start_exon_align in start_exon_align_lines:
        chain = find_exon_chain(sorted_gene_exons, segment_length_table, align_db, 0, start_exon_align)
        if chain:
            chain.append(start_exon_align)
            if not in_one_segment(chain):
                print(chain_to_string(gene_name, reversed(chain)))
                aligned_gene += 1
    return aligned_gene


def find_exon_chain(sorted_gene_exons, segment_length_table, align_db, current_exon_position, current_exon_align):
    position = current_exon_position + 1
    next_exon = sorted_gene_exons[position]
    exon_distance = calc_exon_distance(sorted_gene_exons[current_exon_position], next_exon)
    for exon_align in sorted(align_db[next_exon.name]):
        if not check_pair_exon_alignment(segment_length_table, current_exon_align, exon_align, exon_distance):
            continue
        if position + 1 == len(sorted_gene_exons):
            return [exon_align]
        chain = find_exon_chain(sorted_gene_exons, segment_length_table,
                                align_db, position, exon_align)
        if chain:
            chain.append(exon_align)
            return chain
    else:
        return []


def check_pair_exon_alignment(segment_length_table, align1, align2, exon_distance):
    if align1.strand != align2.strand:
        return False

    if calc_aligned_exons_distance(segment_length_table, align1, align2) > \
            exon_distance * EXON_DISTANCE_INCREASE_MULTIPLIER:
        return False

    if align2.segment_id == align1.segment_id and align1.start > align2.start:
        return False
    return True


def in_one_segment(chain):
    return all(ring[0] == chain[0][0] for ring in chain)


def chain_to_string(gene_name, chain):
    return '{}\t{}'.format(gene_name, ' -> '.join('{}({})'.format(align.segment_id, align.start) for align in chain))


def calc_aligned_exons_distance(segment_length_table, align1, align2):
    if align1.segment_id == align2.segment_id:
        return align2.start - align1.end
    # todo debug concatenation: remove segm
    return segment_length_table['segm'+str(align1.segment_id)] - align1.end + align2.start


def calc_exon_distance(exon1, exon2):
    return exon2.start + exon1.minus_end


if __name__ == '__main__':
    main()